(* module which implements a context graph data structure, including it's supported operations *) 
module CG
    open System.Collections.Generic
    open BJK
    open BIT_ARRAY
    (* this "global" variable stores the number of hash table buckets in graphs created through the newGraph() function.  
        in order to improve performance - it should be initialized to at least 2^X, where X is the number of variables in the analysed program. *)
    let numBuckets = ref 100;
 
    (* a flag indicating if edges with empty DFRs should be added to graph. *) 
    let keepEmptyEdges = ref false;
    
    let debug = ref false;
          
    let debugPrint(s) = if (!debug) then (printfn "%s" s) else ();
   
    type Context = Bitarray

    (* the data which is hanging on the context graph edges *)
    type edgeTag = sop_bar.DFR ref               

    (*the graph type - it is actually a hash table with BitArray keys *)   
    type graph = Dictionary<Context, Dictionary<Context, edgeTag>>

    (* a tagged or untagged edge datatype *)
    type edge = E of Context * Context * edgeTag | UE of Context * Context

    (* returns the specified edge's source context *)    
    let edgeSrc = function
        | E(p,_,_) -> p
        | UE(p,_)  -> p

    (* returns the specified edge's target context *)          
    let edgeTgt = function
        | E(_,q,_) -> q
        | UE(_,q)  -> q
  
    (* If dfr is non-empty:
       1. edge exists in g - updates exiting edge's DFR to be: LUB(existing_DFR,  new_DFR) .
       2. Otherwise, adds the new edge to g. *)
    let addEdge(g : graph, e, cloneDfrs) =
        match e with
        | E(p,q,dfr) ->
                if ( !keepEmptyEdges || (not (sop_bar.dfrEmpty(!dfr)))) then
                    if (g.ContainsKey(p)) then
                        if g.[p].ContainsKey(q) then 
                            let dfr' = if cloneDfrs then ref !(g.[p].[q]) else g.[p].[q]
                            g.[p].[q] := sop_bar.sqcup(!dfr',!dfr)
                            (* in case edge exist in graph update edge to be LUB(existing_DFR, new_DFR) *)
                        else
                            g.[p].Add(q, if cloneDfrs then ref !dfr else dfr)
                            (* Otherwise - insert new edge to g. *)
                    else g.Add(p,new Dictionary<Context,edgeTag>())
                         g.[p].Add(q, if cloneDfrs then ref !dfr else dfr)
                else ()
        | _ -> Util.fail("CG.addEdge: invalid argument")

    (* inserts a list of edges to the graph g. 
        If a specified edge is already in g, the new edge's DFR will be joined with the existing edge's DFR *)
    let rec addEdges = function
        | (g : graph, [],_) -> ()
        | (g : graph, e::el,cloneDfrs) -> (addEdge(g,e,cloneDfrs); addEdges(g,el,cloneDfrs));

    let removeEdge(g : graph, e : edge) : bool =
        match e with
        | UE(p,q) ->
            if g.ContainsKey(p) && g.[p].ContainsKey(q) then
                let _ = g.[p].Remove(q)
                let _ = if g.[p].Count = 0 then g.Remove(p) else false
                true 
            else false
        | _ -> Util.fail("CG.removeEdge: invalid argument")  

    let removeEdgesByPreContext(g : graph ref, pc : Context) = (!g).Remove(pc)

             
      (* constructor for new graphs - creates a new graph given the specified edge list *)
    let newGraph = function
        | [] -> new graph()
        | edgeList ->  
            let g = new graph()
            addEdges(g,edgeList,false)
            g

    let newGraphWithClonedDFRs = function
        | [] -> new graph()
        | edgeList ->  
            let g = new graph()
            addEdges(g,edgeList,true)
            g
                
    (* returns the edge tag of the specified edge *)
    let getEdgeTag(g : graph, e) = 
        match e with
        | UE(p,q) -> if g.ContainsKey(p) && g.[p].ContainsKey(q) then
                                            Some g.[p].[q]
                                          else
                                            None
        | _ -> Util.fail("CG.getEdgeTag: invalid argument")

    (* returns a list of all edges in graph - not including edge tags (UE edge constructor is used to create each edge) *)    
    let untaggedEdges(g : graph) =
        let pUntaggedEdges(p : Context,tgt : Dictionary<Context,edgeTag>) =
            let tgtContexts = (Seq.toList tgt.Keys)
            List.map (fun (q : Context) -> UE(p,q)) tgtContexts      
        let ps = List.map (fun c -> (c,g.[c])) (Seq.toList(g.Keys)) 
        List.concat (List.map pUntaggedEdges ps)

    (* returns all edges in graph including edge tags (E edge constructor is usedto create each edge) *)
    let edges(g : graph) = 
        let pTaggedEdges(p : Context,tgt : Dictionary<Context,edgeTag>) =
              let tgtContextAndDfrPairs = (Seq.toList tgt) 
              let pKeyValue2edge(pair : KeyValuePair<Context,edgeTag> ) =
                    E(p, pair.Key, pair.Value)
              List.map pKeyValue2edge tgtContextAndDfrPairs
               
        let ps = List.map (fun c -> (c,g.[c])) (Seq.toList(g.Keys))  
        List.concat (List.map pTaggedEdges ps)

    
    (* creates an empty context which can hold up to the specified number of variables *)            
    let makeEmptyContext(size : int) : Context = Bitarray.create(size)

    (* creates a context which can hold up to the specified number of variables and is filled with the specified list of variable indices *)
    let makeContext(size : int, vars : int list) : Context =
        let c = makeEmptyContext(size)
        List.iter (fun i->(Bitarray.set (i-1) c)) vars
        c

    (* returns a copy of the srcContext, assuming that numVars is the number of variables in srcContext *)
    let copyContext(srcContext) = Bitarray.copy srcContext
    
    (* adds variable with index x to context c *)
    let addVarToContext(x, c : Context) = Bitarray.set (x-1) c

    (* removes variable with index x to context c *)    
    let removeVarFromContext(x, c : Context) = Bitarray.unset (x-1) c
    
    (* if 'add' is true - adds x to context c, otherwise removes x from context c *)
    let setVarInContext(x, c, add) = if (add) then addVarToContext(x,c) else removeVarFromContext(x,c)

    (* returns a list of variable indices which are contained in the specified context. *)
    let contextToVarList(c : Context) : int list = 
        List. map (fun i -> i+1) (Bitarray.setBits c)
     
     (* outputs the context's variables to a nicely formatted string. *)    
    let prtContext(c : Context) = 
        let prtVar(index, comma) = if (Bitarray.isSet index c) then 
                                       (if (comma) then ", " else "") + (index+1).ToString() 
                                   else ""
        
        let rec prtVars(index, len,first) = 
            if (index < len) then  
                  let vr = prtVar(index, not first)
                  let first1 = first && ((String.length vr) = 0)
                  vr + prtVars(index+1, len, first1) 
            else ""                                                 
        
        ("{ " + prtVars(0, (Bitarray.length c), true) + " }")
      

    let rec prtContexts = function
        | []   -> ""
        | c::t -> prtContext(c) + ", " + prtContexts(t)
      
    let inContext(v : int, c : Context) = Bitarray.isSet (v-1) c
    
    let contextToNum(c : Context) = 
        let rec computeNum = function
            | (numSoFar : int, powerOf2 : int, 0) -> numSoFar
            | (numSoFar : int, powerOf2, index : int) -> computeNum(
                                                           (if (Bitarray.isSet index c) then numSoFar+powerOf2 else 0),
                                                           powerOf2 * 2, 
                                                           index-1) 
       
        computeNum(0,1,(Bitarray.length c)-1)
        
    (* a partial order on edges, based on lexicographic ordering of the source and target contexts hash code. *)    
    let edgeSmaller(e1,e2) = 
        let p = contextToNum(edgeSrc(e1))
        let p' = contextToNum(edgeSrc(e2))
        p < p' || (p = p' && (contextToNum(edgeTgt(e1)) < contextToNum(edgeTgt(e2))))

    (* inserts an edge into a sorted edge list *)
    let rec insertEdge = function
        | (e,[]) -> [e]
        | (e,((h::t) as el)) -> if (edgeSmaller(h,e)) then 
                                  h::insertEdge(e,t) 
                                else 
                                  e::el
    (* sort an edge list by lexical order <src_context, tgt_context> *)
    let rec sortEdges = function
        | [] -> []
        | e::t -> insertEdge(e,sortEdges(t));

    (* returns true iff context c1 is equal to context c2. *)        
    let eqContexts(c1 : Context, c2 : Context) = Bitarray.areAllUnset(c1 ^^^ c2)

    let contextSmaller(c1,c2) = 
        let p = contextToNum(c1)
        let p' = contextToNum(c2)
        p < p'
    
    (* inserts a context into a sorted context list *)
    let rec insertContext c cl = 
        match cl with
        | []           -> [c]
        | (h::t) as cl -> 
            if (contextSmaller(h,c)) then 
                h::(insertContext c t)
            else c::cl
      
    let sortContexts(contextList : Context list) = 
        List.sort contextList

        (* returns all pre contexts of graph g. *)
    let preContexts(g : graph) : Context list = 
        List.map (fun e -> match e with 
                           |UE(p,_)-> p
                           | _ -> Util.fail("CG.preContexts: Invalid argument."))
                 (untaggedEdges(g))

    (* returns true if every context of preContextsList is a pre-contexts of graph g and vice-versa. *)    
    let preContextsEqual(preContextsList, g) = 
        let rec equalContextLists = function
            | ([],[]) -> true
            | (_,[]) -> false
            | ([],_) -> false
            | (c1::cl1, c2::cl2) -> eqContexts(c1,c2) && equalContextLists(cl1,cl2)
        equalContextLists(preContexts(g),preContextsList)
     
    (* returns true if every context of preContextsList is a pre-contexts of graph g. *)
    let preContextsIn(preContextsLis, g : graph) =
        let yes = ref true;
        for c in preContextsLis do
            yes := !yes && g.ContainsKey(c)
        !yes        
        
    let prtDfr((sop_bar.Dfr(M,R)) as d) = "\n    M=[ " + mat.prt(M) + "]\n    R=[ " + sop_bar.prtR(d) + "]";

    (* outputs the edge context variables and DFR to a nicely formatted string. *)          
    let prtEdge = function
        | E(p, q, dfr) -> match !dfr with
                          | (sop_bar.Dfr(M, R)) as d -> prtContext(p) + "-->" + prtContext(q) + " " + prtDfr(d)
        | UE(p, q)     -> prtContext(p) + " ? "  + prtContext(q);

    (* for all edge in the list - outputs the edge context variables and DFR to a nicely formatted string. *)          
    let rec prtEdgeList = function
        | ([]) -> ""
        | (e::t) -> prtEdge(e) + "\n" + prtEdgeList(t);

     (* for all edge in the graph - outputs the edge context variables and DFR to a nicely formatted string. *)    
    let prtEdges(g : graph) =
        prtEdgeList(sortEdges(edges(g)))

    (* for all edge in the graph - outputs the edge context variables and DFR to a nicely formatted string. *)    
    let printEdges(g : graph) =
        printfn "%s" (prtEdgeList(sortEdges(edges(g))))
    
    (* returns all post contexts of graph g. *)
    let postContexts(g : graph) : Context list = List.map edgeTgt (untaggedEdges(g));

    let clone(g : graph) : graph = 
        newGraphWithClonedDFRs(edges(g));

    (* returns true if g1 and g2 are equal in all aspects *)
    let equal(g1 : graph, g2 : graph) : bool =
        let res = ref true
        let mutable e = g1.Keys.GetEnumerator()
        let () = while e.MoveNext() && !res do
                    let p = e.Current
                    if (g2.ContainsKey(p)) then
                        let qds = g1.[p]
                        let qds' = g2.[p]
                        let mutable e1 = qds.Keys.GetEnumerator()
                        while e1.MoveNext() && !res do
                            let q = e1.Current
                            if qds'.ContainsKey(q) then
                                res := !res && (BJK.sop_bar.dfrsame(!(qds.[q]),!(qds'.[q])))     
                            else res := false 
                    else 
                        res := false
        !res
        
    (* returns true if c is a context which contains all variables (all bits are 1). *)
    let isFullContext(c : Context) = Bitarray.areAllSet c
        
    
    (* computes all possible contexts (except the full context) over a set of variables, 
       where highestVar is the highest variable index *) 
    let allContexts(highestVar) =
        let rec intToBinaryInternal = function 
            | (0,n,k) -> []
            | (i,n,k) -> if ((i % 2) = 1) then 
                            (n-k+1)::intToBinaryInternal(i / 2,n,k-1) 
                         else 
                            intToBinaryInternal(i / 2,n,k-1) ;
        let intToBinaryList(n) =
            intToBinaryInternal(n,n,n);
            
        let rec makeAllContexts(highest_var,i) = 
            let c = makeContext(highest_var, intToBinaryList(i));
            if (isFullContext(c)) then [] else  c :: makeAllContexts(highest_var,i+1)
            
        
        let l = makeAllContexts(highestVar,0);
        let _ = debugPrint("All contexts: " + prtContexts(l) + "\n"); 
        l
        
    
        
    (* compute the composition of two graphs *)
    let compose(g1 : graph, g2 : graph) =                        
        let g = newGraph([])
        let () = for p in g1.Keys do
                    for qd in g1.[p] do
                        let found, qTagDfrs = g2.TryGetValue(qd.Key)
                        if (found) then
                            let dfr = qd.Value 
                            for qTagDfr in qTagDfrs do
                                let q' = qTagDfr.Key
                                let dfr' = qTagDfr.Value              
                                addEdge(g, E(p,q',ref (sop_bar.prod(!dfr,!dfr'))),false) 
                        else ()
        g
    
    let debugPrintEdges(edges') =
        if (!debug) then debugPrint(prtEdgeList(edges')) else ()
        
    (* computes the join (LUB) of two graphs *)    
    let join(g1 : graph, g2 : graph) : graph= 
        newGraph( ((edges(g1)) @ (edges(g2))) )
 
    (* computes the join (LUB) of two graphs, cloning DFRs on the way. *)    
    let joinWithClonedDFRs(g1 : graph, g2 : graph) : graph= 
        newGraphWithClonedDFRs( ((edges(g1)) @ (edges(g2))) )
   
    (* removes from g all edges which do not exit a context from pc. *)
    let filterEdgesByPreContexts(g : graph ref, pc : Context list) =
        let pcs = Set.ofList pc
        let pcg = preContexts(!g)
        for p in pcg do
            if not (pcs.Contains p) then
                let _ = removeEdgesByPreContext(g, p)
                ()
            else ()
                             
 (* module CG *)

